{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 在 MLFlow 中使用自定义 Python 封装"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本 notebook 展示如何在 Phi-3 mini 4K Instruct 示例中使用自定义 Python 封装。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入所需的库\n",
    "import mlflow\n",
    "from mlflow.models import infer_signature\n",
    "import onnxruntime_genai as og"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 为 Phi-3 Mini 4K 模型定义 Python 类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 自定义 Python 类\n",
    "class Phi3Model(mlflow.pyfunc.PythonModel):\n",
    "    def load_context(self, context):\n",
    "        # 从工件中检索模型\n",
    "        model_path = context.artifacts[\"phi3-mini-onnx\"]\n",
    "        model_options = {\n",
    "             \"max_length\": 300,\n",
    "             \"temperature\": 0.2,         \n",
    "        }\n",
    "    \n",
    "        # 定义模型\n",
    "        self.phi3_model = og.Model(model_path)\n",
    "        self.params = og.GeneratorParams(self.phi3_model)\n",
    "        self.params.set_search_options(**model_options)\n",
    "        \n",
    "        # 定义分词器（Tokenizer）\n",
    "        self.tokenizer = og.Tokenizer(self.phi3_model)\n",
    "\n",
    "    def predict(self, context, model_input):\n",
    "        # 从输入中检索提示\n",
    "        prompt = model_input[\"prompt\"][0]\n",
    "        self.params.input_ids = self.tokenizer.encode(prompt)\n",
    "\n",
    "        # 生成模型的响应\n",
    "        response = self.phi3_model.generate(self.params)\n",
    "\n",
    "        return self.tokenizer.decode(response[0][len(self.params.input_ids):])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成 MLFlow 工件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/06/18 14:17:34 INFO mlflow.models.utils: We convert input dictionaries to pandas DataFrames such that each key represents a column, collectively constituting a single row of data. If you would like to save data as multiple rows, please convert your data to a pandas DataFrame before passing to input_example.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c43e3ebd0d74c18997ec6bddf560e47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 使用自定义 Python 模型生成 MLflow 模型\n",
    "input_example = {\"prompt\": \"<|system|>You are a helpful AI assistant.<|end|><|user|>What is the capital of Spain?<|end|><|assistant|>\"}\n",
    "artifact_path = \"phi3_mlflow_model\"\n",
    "\n",
    "with mlflow.start_run() as run:\n",
    "    model_info = mlflow.pyfunc.log_model(\n",
    "        artifact_path = artifact_path,\n",
    "        python_model = Phi3Model(),\n",
    "        artifacts = {\n",
    "            \"phi3-mini-onnx\": \"cpu_and_mobile/cpu-int4-rtn-block-32-acc-level-4\",\n",
    "        },\n",
    "        input_example = input_example,\n",
    "        signature = infer_signature(input_example, [\"Run\"]),\n",
    "        extra_pip_requirements = [\"torch\", \"onnxruntime_genai\", \"numpy\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running Phi-3 as MLFlow model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载 Phi-3 MLFlow 模型\n",
    "loaded_model = mlflow.pyfunc.load_model(\n",
    "    model_uri = model_info.model_uri\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "inputs: \n",
       "  ['prompt': string (required)]\n",
       "outputs: \n",
       "  [string (required)]\n",
       "params: \n",
       "  None"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 检查 MLFlow 模型的签名\n",
    "loaded_model.metadata.signature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Alright, here's a little atom-related joke for you!\n",
      "\n",
      "Why don't electrons ever play hide and seek with protons?\n",
      "\n",
      "Because good luck finding them when they're always \"sharing\" their electrons!\n",
      "\n",
      "Remember, this is all in good fun, and we're just having a little atomic-level humor!\n"
     ]
    }
   ],
   "source": [
    "# 测试加载的模型\n",
    "response = loaded_model.predict(\n",
    "    {\"prompt\": \"<|system|>You are a stand-up comedian.<|end|><|user|>Tell me a joke about atom<|end|><|assistant|>\",}\n",
    ")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
